from typing import List, Optional

from pydantic import BaseModel


class ChatMessage(BaseModel):
    role: Optional[str] = "human"
    content: Optional[str] = ""


class Role(BaseModel):
    name: Optional[str] = ""
    age: Optional[str] = ""
    sex: Optional[str] = ""
    company: Optional[str] = ""
    company_abbreviation: Optional[str] = ""
    brand: Optional[str] = ""
    industry: Optional[str] = "4S店"
    position: Optional[str] = "专业的销售客服"
    characters: Optional[str] = "在回复客户问题时，请使用亲切且热情的口吻，可以称呼客户为‘老板’，并适当多使用语气助词来增强交流的亲切感。"
    company_province: Optional[str] = ""
    company_city: Optional[str] = ""
    company_address: Optional[str] = ""
    company_tel: Optional[str] = ""
    company_opening: Optional[str] = ""
    company_liveTime: Optional[str] = ""
    company_liveToStore_discount: Optional[str] = ""
    company_testDrive_discount: Optional[str] = ""
    company_vehicleSeries: Optional[str] = ""
    company_other: Optional[str] = ""


class Customer(BaseModel):
    """客户信息"""
    # 车型
    custom_vehicle_type: Optional[str] = ""
    custom_phone_number: Optional[str] = ""
    custom_weChat_account: Optional[str] = ""
    custom_province: Optional[str] = ""
    custom_city: Optional[str] = ""
    custom_area: Optional[str] = ""


class BaseRequest(BaseModel):
    request_id: str = ""
    model_name: str = "gpt-3.5-turbo-16k"
    prompt: str
    chat_history: List[ChatMessage] = []
    stream: Optional[bool] = True
    temperature: float = 0.5
    top_p: float = 0.0
    max_tokens: int = 256
    vector_search_distance: float = 0.5
    vector_search_limit: int = 10


class CompletionRequest(BaseRequest):
    model_style: str = "stable"
    vector_categories: List[str] = []
    role: Role = Role()
    custom: Customer = Customer()


class ExtractionRequest(BaseRequest):
    vehicle_brand_ids: List[str] = []
